/*
 * FluidSolver.java
 * Alexander McKenzie
 * 12 March, 2004
 *
 */

/**
 * Jos Stam style fluid solver with vorticity confinement
 * and buoyancy force.
 *
 * @author Alexander McKenzie
 * @version 1.0
 **/

public class FluidSolver
{
    int n, size;
    float dt;

    float visc = 0.0f;
    float diff = 0.0f;

    float[][] tmp;

    float[][] d, dOld;
    float[][] u, uOld;
    float[][] v, vOld;
    float[][] curl;
    float[][] temp;


    /**
     * Set the grid size and timestep.
     **/

    public void setup(int n, float dt)
    {
        this.n = n;
        this.dt = dt;
        size = (n + 2);

        reset();
    }


    /**
     * Reset the datastructures.
     * We use 1d arrays for speed.
     **/

    public void reset()
    {
        d    = new float[size][size];
        dOld = new float[size][size];
        u    = new float[size][size];
        uOld = new float[size][size];
        v    = new float[size][size];
        vOld = new float[size][size];
        curl = new float[size][size];
        temp = new float[size][size];

        for (int i = 0; i < size; i++)
        {
        	for (int j = 0; j < size; j++)
        	{
        		u[i][j] = uOld[i][j] = v[i][j] = vOld[i][j] = 0.0f;
        		d[i][j] = dOld[i][j] = curl[i][j] = 0.0f;
        		temp[i][j] = 0.0f;
        	}
        }
        
        d[size / 2][size - 2] = 100;
        d[size / 2][1] = 100;
    }


    /**
     * Calculate the buoyancy force as part of the velocity solver.
     * Fbuoy = -a*d*Y + b*(T-Tamb)*Y where Y = (0,1). The constants
     * a and b are positive with appropriate (physically meaningful)
     * units. T is the temperature at the current cell, Tamb is the
     * average temperature of the fluid grid. The density d provides
     * a mass that counteracts the buoyancy force.
     *
     * In this simplified implementation, we say that the tempterature
     * is synonymous with density (since smoke is *hot*) and because
     * there are no other heat sources we can just use the density
     * field instead of a new, seperate temperature field.
     *
     * @param Fbuoy Array to store buoyancy force for each cell.
     **/

    public void buoyancy(float[][] Fbuoy)
    {
        float Tamb = 0;
        float a = 0.000625f;
        float b = 0.025f;

        // sum all temperatures
        for (int i = 1; i <= n; i++)
        {
            for (int j = 1; j <= n; j++)
            {
                Tamb += d[i][j];
            }
        }

        // get average temperature
        Tamb /= (n * n);

        // for each cell compute buoyancy force
        for (int i = 1; i <= n; i++)
        {
            for (int j = 1; j <= n; j++)
            {
                Fbuoy[i][j] = a * d[i][j] + -b * (d[i][j] - Tamb);
            }
        }
    }


    /**
     * Calculate the curl at position (i, j) in the fluid grid.
     * Physically this represents the vortex strength at the
     * cell. Computed as follows: w = (del x U) where U is the
     * velocity vector at (i, j).
     *
     * @param i The x index of the cell.
     * @param j The y index of the cell.
     **/

    public float curl(int i, int j)
    {
        float du_dy = (u[i][j + 1] - u[i][j - 1]) * .5f;
        float dv_dx = (v[i + 1][j] - v[i - 1][j]) * .5f;

        return dv_dx - du_dy;
    }


    /**
     * Calculate the vorticity confinement force for each cell
     * in the fluid grid. At a point (i,j), Fvc = N x w where
     * w is the curl at (i,j) and N = del |w| / |del |w||.
     * N is the vector pointing to the vortex center, hence we
     * add force perpendicular to N.
     *
     * @param Fvc_x The array to store the x component of the
     *        vorticity confinement force for each cell.
     * @param Fvc_y The array to store the y component of the
     *        vorticity confinement force for each cell.
     **/

    public void vorticityConfinement(float[][] Fvc_x, float[][] Fvc_y)
    {
        float dw_dx, dw_dy;
        float length;
        float v;

        // Calculate magnitude of curl(u,v) for each cell. (|w|)
        for (int i = 1; i <= n; i++)
        {
            for (int j = 1; j <= n; j++)
            {
                curl[i][j] = Math.abs(curl(i, j));
            }
        }

        for (int i = 2; i < n; i++)
        {
            for (int j = 2; j < n; j++)
            {

                // Find derivative of the magnitude (n = del |w|)
                dw_dx = (curl[i + 1][j] - curl[i - 1][j]) * 0.5f;
                dw_dy = (curl[i][j + 1] - curl[i][j - 1]) * 0.5f;

                // Calculate vector length. (|n|)
                // Add small factor to prevent divide by zeros.
                length = (float) Math.sqrt(dw_dx * dw_dx + dw_dy * dw_dy)
                         + 0.000001f;

                // N = ( n/|n| )
                dw_dx /= length;
                dw_dy /= length;

                v = curl(i, j);

                // N x w
                Fvc_x[i][j] = dw_dy * v;
                Fvc_y[i][j] = dw_dx * -v;
            }
        }
    }


    /**
     * The basic velocity solving routine as described by Stam.
     **/

    public void velocitySolver()
    {

        // add velocity that was input by mouse
        addSource(u, uOld);
        addSource(v, vOld);

        // add in vorticity confinement force
        vorticityConfinement(uOld, vOld);
        addSource(u, uOld);
        addSource(v, vOld);

        // add in buoyancy force
        buoyancy(vOld);
        addSource(v, vOld);

        // swapping arrays for economical mem use
        // and calculating diffusion in velocity.
        swapU();
        diffuse(0, u, uOld, visc);

        swapV();
        diffuse(0, v, vOld, visc);

        // we create an incompressible field
        // for more effective advection.
        project(u, v, temp, vOld);

        swapU(); swapV();

        // self advect velocities
        advect(1, u, uOld, uOld, vOld);
        advect(2, v, vOld, uOld, vOld);

        // make an incompressible field
        project(u, v, uOld, vOld);

        // clear all input velocities for next frame
        for (int i = 0; i < size; i++)
        { 
        	for (int j = 0; j < size; j++)
        	{
        		uOld[i][j] = 0; 
        		vOld[i][j] = 0;
        	}
        }
    }


    /**
     * The basic density solving routine.
     **/

    public void densitySolver()
    {
        // add density inputted by mouse
        addSource(d, dOld);
        swapD();

        diffuse(0, d, dOld, diff);
        swapD();

        advect(0, d, dOld, u, v);

        // clear input density array for next frame
        for (int i = 0; i < size; i++) 
        	for (int j = 0; j < size; j++)
        		dOld[i][j] = 0;
    }


    private void addSource(float[][] x, float[][] x0)
    {
        for (int i = 0; i < size; i++)
        	for (int j = 0; j < size; j++)
        		x[i][j] += dt * x0[i][j];
    }


    /**
     * Calculate the input array after advection. We start with an
     * input array from the previous timestep and an and output array.
     * For all grid cells we need to calculate for the next timestep,
     * we trace the cell's center position backwards through the
     * velocity field. Then we interpolate from the grid of the previous
     * timestep and assign this value to the current grid cell.
     *
     * @param b Flag specifying how to handle boundries.
     * @param d Array to store the advected field.
     * @param d0 The array to advect.
     * @param du The x component of the velocity field.
     * @param dv The y component of the velocity field.
     **/

    private void advect(int b, float[][] d, float[][] d0, float[][] du, float[][] dv)
    {
        int i0, j0, i1, j1;
        float x, y, s0, t0, s1, t1, dt0;

        dt0 = dt * n;

        for (int i = 1; i <= n; i++)
        {
            for (int j = 1; j <= n; j++)
            {
                // go backwards through velocity field
                x = i - dt0	* du[i][j];
                y = j - dt0 * dv[i][j];

                // interpolate results
                if (x > n + 0.5) x = n + 0.5f;
                if (x < 0.5)     x = 0.5f;

                i0 = (int) x;
                i1 = i0 + 1;

                if (y > n + 0.5) y = n + 0.5f;
                if (y < 0.5)     y = 0.5f;

                j0 = (int) y;
                j1 = j0 + 1;

                s1 = x - i0;
                s0 = 1 - s1;
                t1 = y - j0;
                t0 = 1 - t1;

                //place modified
                //d[i][j] = (d0[i0][j0] + d0[i0][j1] + d0[i1][j0] + d0[i1][j1]) / 4.0f;
                d[i][j] = s0 * (t0 * d0[i0][j0] + t1 * d0[i0][j1])
                        + s1 * (t0 * d0[i1][j0] + t1 * d0[i1][j1]);

            }
        }
        setBoundry(b, d);
    }



    /**
     * Recalculate the input array with diffusion effects.
     * Here we consider a stable method of diffusion by
     * finding the densities, which when diffused backward
     * in time yield the same densities we started with.
     * This is achieved through use of a linear solver to
     * solve the sparse matrix built from this linear system.
     *
     * @param b Flag to specify how boundries should be handled.
     * @param c The array to store the results of the diffusion
     * computation.
     * @param c0 The input array on which we should compute
     * diffusion.
     * @param diff The factor of diffusion.
     **/

    private void diffuse(int b, float[][] c, float[][] c0, float diff)
    {
        float a = dt * diff * n * n;
        linearSolver(b, c, c0, a, 1 + 4 * a);
    }


    /**
     * Use project() to make the velocity a mass conserving,
     * incompressible field. Achieved through a Hodge
     * decomposition. First we calculate the divergence field
     * of our velocity using the mean finite differnce approach,
     * and apply the linear solver to compute the Poisson
     * equation and obtain a "height" field. Now we subtract
     * the gradient of this field to obtain our mass conserving
     * velocity field.
     *
     * @param x The array in which the x component of our final
     * velocity field is stored.
     * @param y The array in which the y component of our final
     * velocity field is stored.
     * @param p A temporary array we can use in the computation.
     * @param div Another temporary array we use to hold the
     * velocity divergence field.
     *
     **/

    void project(float[][] x, float[][] y, float[][] p, float[][] div)
    {
        for (int i = 1; i <= n; i++)
        {
            for (int j = 1; j <= n; j++)
            {
                div[i][j] = (x[i + 1][j] - x[i - 1][j]
                              + y[i][j + 1] - y[i][j - 1])
                              * - 0.5f / n;
                p[i][j] = 0;
            }
        }

        setBoundry(0, div);
        setBoundry(0, p);

        linearSolver(0, p, div, 1, 4);

        for (int i = 1; i <= n; i++)
        {
            for (int j = 1; j <= n; j++)
            {
                x[i][j] -= 0.5f * n * (p[i + 1][j] - p[i - 1][j]);
                y[i][j] -= 0.5f * n * (p[i][j + 1] - p[i][j - 1]);
            }
        }

        setBoundry(1, x);
        setBoundry(2, y);
    }


    /**
     * Iterative linear system solver using the Gauss-sidel
     * relaxation technique. Room for much improvement here...
     *
     **/

    void linearSolver(int b, float[][] x, float[][] x0, float a, float c)
    {
        for (int k = 0; k < 20; k++)
        {
            for (int i = 1; i <= n; i++)
            {
                for (int j = 1; j <= n; j++)
                {
                    x[i][j] = (a * ( x[i - 1][j] + x[i + 1][j]
                                    +   x[i][j  - 1] + x[i][j + 1])
                                    +  x0[i][j]) / c;
                }
            }
            setBoundry(b, x);
        }
    }


    // specifies simple boundry conditions.
    private void setBoundry(int b, float[][] x)
    {
        for (int i = 1; i <= n; i++)
        {
            x[0][i] = b == 1 ? -x[1][i] : x[1][i];
            x[n + 1][i] = b == 1 ? -x[n][i] : x[n][i];
            x[i][0] = b == 2 ? -x[i][1] : x[i][1];
            x[i][n + 1] = b == 2 ? -x[i][n] : x[i][n];
        }

        x[0][0] = 0.5f * (x[1][0] + x[0][1]);
        x[0][n + 1] = 0.5f * (x[1][n + 1] + x[0][n]);
        x[n + 1][0] = 0.5f * (x[n][0] + x[n + 1][1]);
        x[n + 1][n+1] = 0.5f * (x[n][n + 1] + x[n + 1][n]);
    }

    // util array swapping methods
    public void swapU(){ tmp = u; u = uOld; uOld = tmp; }
    public void swapV(){ tmp = v; v = vOld; vOld = tmp; }
    public void swapD(){ tmp = d; d = dOld; dOld = tmp; }
}
